import numpy as np
import os
import random
import torch
import copy
import torch.nn as nn
from lib.utils.tools import read_pkl
from lib.utils.utils_data import flip_data, crop_scale_3d
    
class Augmenter2D(object):
    """
        Make 2D augmentations（增强） on the fly. PyTorch batch-processing GPU version.
    """
    def __init__(self, args):
        self.d2c_params = read_pkl(args.d2c_params_path)
        self.noise = torch.load(args.noise_path)
        self.mask_ratio = args.mask_ratio
        self.mask_T_ratio = args.mask_T_ratio
        self.num_Kframes = 27
        self.noise_std = 0.002

    def dis2conf(self, dis, a, b, m, s):
        f = a/(dis+a)+b*dis
        shift = torch.randn(*dis.shape)*s + m
        # if torch.cuda.is_available():
        shift = shift.to(dis.device)
        return f + shift
    
    def add_noise(self, motion_2d):
        """
        这段代码是一个深度学习模型中的一个函数，它的作用是向输入的二维运动数据（motion_2d）添加噪声。这个过程可能是为了数据增强、模型鲁棒性提升或者模拟不确定性等目的。下面是对代码的逐行解读：

        首先，函数从类的成员变量d2c_params和noise中提取了一些参数，包括a, b, m, s以及噪声的均值（mean）和标准差（std）。

        接着，代码检查是否存在一个名为uniform_range的键在噪声参数字典中，如果不存在，则将其默认值设为0.06。

        然后，代码提取了输入数据motion_2d的形状信息，包括批量大小（batch_size）、帧数（num_frames）和关节点数量（num_joints）。

        接下来，代码生成了一些随机样本和高斯样本。sel是一个随机采样的张量，用于后续选择噪声类型。gaussian_sample是一个高斯分布的噪声样本，uniform_sample是一个均匀分布的噪声样本。

        delta_noise是一个额外的高斯噪声，用于在整个序列上添加噪声。

        代码将所有的噪声样本和参数转移到了与motion_2d相同的设备上（可能是GPU或CPU）。

        接下来，代码根据sel张量的值选择使用高斯噪声还是均匀噪声，并将它们组合成delta。

        delta_expand通过插值将delta扩展到原始的运动数据帧数。

        delta_final是最终的噪声，它由delta_expand和delta_noise相加得到。

        motion_2d加上delta_final得到添加了噪声的运动数据。

        dx和dy分别是添加噪声后的运动数据在x和y方向上的分量。

        dis2和dis分别计算了每个关节点的位移的平方和开方，用于后续的计算。

        conf是根据位移计算得到的信心度（confidence），它是通过调用dis2conf函数并使用a, b, m, s作为参数得到的。这个信心度可能用于表示模型对预测的置信程度。

        最后，函数返回一个新的张量，它将添加噪声后的运动数据和信心度拼接在一起。

        整体来看，这段代码实现了一个向二维运动数据添加噪声的过程，并通过计算信心度来评估模型对每个预测的置信程度。这可能是在一个生成模型或者动作识别模型中用于提高模型鲁棒性和泛化能力的一种技术。
        """
        a, b, m, s = self.d2c_params["a"], self.d2c_params["b"], self.d2c_params["m"], self.d2c_params["s"]
        if "uniform_range" in self.noise.keys():
            uniform_range = self.noise["uniform_range"]
        else:
            uniform_range = 0.06
        motion_2d = motion_2d[:,:,:,:2]
        batch_size = motion_2d.shape[0]
        num_frames = motion_2d.shape[1]
        num_joints = motion_2d.shape[2]
        mean = self.noise['mean'].float()
        std = self.noise['std'].float()
        weight = self.noise['weight'][:,None].float()
        sel = torch.rand((batch_size, self.num_Kframes, num_joints, 1))
        gaussian_sample = (torch.randn(batch_size, self.num_Kframes, num_joints, 2) * std + mean) 
        uniform_sample = (torch.rand((batch_size, self.num_Kframes, num_joints, 2))-0.5) * uniform_range
        noise_mean = 0
        delta_noise = torch.randn(num_frames, num_joints, 2) * self.noise_std + noise_mean
        
        # if torch.cuda.is_available():
        mean = mean.to(motion_2d.device)
        std = std.to(motion_2d.device)
        weight = weight.to(motion_2d.device)
        gaussian_sample = gaussian_sample.to(motion_2d.device)
        uniform_sample = uniform_sample.to(motion_2d.device)
        sel = sel.to(motion_2d.device)
        delta_noise = delta_noise.to(motion_2d.device)
            
        delta = gaussian_sample*(sel<weight) + uniform_sample*(sel>=weight)
        delta_expand = torch.nn.functional.interpolate(delta.unsqueeze(1), [num_frames, num_joints, 2], mode='trilinear', align_corners=True)[:,0]
        delta_final = delta_expand + delta_noise      
        motion_2d = motion_2d + delta_final 
        dx = delta_final[:,:,:,0]
        dy = delta_final[:,:,:,1]
        dis2 = dx*dx+dy*dy
        dis = torch.sqrt(dis2)
        conf = self.dis2conf(dis, a, b, m, s).clip(0,1).reshape([batch_size, num_frames, num_joints, -1])
        return torch.cat((motion_2d, conf), dim=3)
        
    def add_mask(self, x):
        ''' motion_2d: (N,T,17,3)
        '''
        N,T,J,C = x.shape
        mask = torch.rand(N,T,J,1, dtype=x.dtype, device=x.device) > self.mask_ratio
        mask_T = torch.rand(1,T,1,1, dtype=x.dtype, device=x.device) > self.mask_T_ratio
        x = x * mask * mask_T
        return x
    
    def augment2D(self, motion_2d, mask=False, noise=False):     
        if noise:
            motion_2d = self.add_noise(motion_2d)
        if mask:
            motion_2d = self.add_mask(motion_2d)
        return motion_2d
    
class Augmenter3D(object):
    """
        Make 3D augmentations when dataloaders get items. NumPy single motion version.
    """
    def __init__(self, args):
        self.flip = args.flip
        if hasattr(args, "scale_range_pretrain"):
            self.scale_range_pretrain = args.scale_range_pretrain
        else:
            self.scale_range_pretrain = None
    
    def augment3D(self, motion_3d):
        if self.scale_range_pretrain:
            motion_3d = crop_scale_3d(motion_3d, self.scale_range_pretrain)
        if self.flip and random.random()>0.5:                       
            motion_3d = flip_data(motion_3d)
        return motion_3d